from xml.sax import make_parser
from xml.sax.handler import ContentHandler


class BoxHandler(ContentHandler):
    def __init__(self):
        super().__init__()
        self.boxes = []
        self.names = ['xmin', 'xmax', 'ymin', 'ymax']
        self.name = ''
        self.box = [0, 0, 0, 0]  # 'xmin', 'ymin', 'xmax', 'ymax'
        self.tag2idx = {'xmin': 0, 'ymin': 1, 'xmax': 2, 'ymax': 3}

    def startDocument(self):
        self.boxes.clear()

    def startElement(self, name, attrs):
        if name in self.names:
            self.name = name

    def characters(self, content):
        if not self.name:
            return
        content = int(content)
        self.box[self.tag2idx[self.name]] = content
        self.name = ''

    def endElement(self, name):
        if name == 'bndbox':
            self.boxes.append(self.box.copy())
            self.name = ''


if __name__ == '__main__':
    parser = make_parser()
    handler = BoxHandler()
    parser.setContentHandler(handler)
    path = r'F:\FlyAI\TBDetection_FlyAI\data\input\TBDetection\xml\15.xml'
    parser.parse(path)
    print(handler.boxes, len(handler.boxes))
